package flink.dataset;

import org.apache.flink.api.common.functions.FlatMapFunction;
import org.apache.flink.api.common.functions.JoinFunction;
import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.functions.MapPartitionFunction;
import org.apache.flink.api.common.operators.Order;
import org.apache.flink.api.java.DataSet;
import org.apache.flink.api.java.ExecutionEnvironment;
import org.apache.flink.api.java.operators.CrossOperator;
import org.apache.flink.api.java.operators.DataSource;
import org.apache.flink.api.java.operators.GroupReduceOperator;
import org.apache.flink.api.java.operators.UnionOperator;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.api.java.tuple.Tuple3;
import org.apache.flink.util.Collector;
import org.junit.Test;

import java.util.Arrays;
import java.util.Iterator;
import java.util.List;

/**
 * @author binarylei
 * @version 2019-11-29
 */
public class OperatorDataSetTest {

    @Test
    public void testMapPartition() throws Exception {
        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();

        DataSet<String> text = env.fromElements(
                "Who's there?",
                "I think I hear them. Stand, ho! Who's there?");

        DataSet<Tuple2<String, Integer>> wordCounts = text
                .mapPartition(new MapPartitionFunction<String, Tuple2<String, Integer>>() {
                    @Override
                    public void mapPartition(Iterable<String> iterable,
                            Collector<Tuple2<String, Integer>> collector) throws Exception {
                        Iterator<String> iterator = iterable.iterator();
                        while (iterator.hasNext()) {
                            String line = iterator.next();
                            String[] arr = line.split("\\s+");
                            Arrays.stream(arr).forEach(word -> collector.collect(new Tuple2<>(word, 1)));
                        }
                    }
                }).groupBy(0).sum(1);

        wordCounts.print();
    }

    @Test
    public void testFlatMapFunction() throws Exception {
        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();

        DataSet<String> text = env.fromElements(
                "Who's there?",
                "I think I hear them. Stand, ho! Who's there?");

        text.flatMap(new FlatMapFunction<String, String>() {
            @Override
            public void flatMap(String line, Collector<String> collector) throws Exception {
                String[] arr = line.split("\\s+");
                Arrays.stream(arr).forEach(word -> collector.collect(word));
            }
        }).distinct().print();
    }

    @Test
    public void testJoin() throws Exception {
        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();

        // <id, name>
        DataSource<Tuple2<Long, String>> names = env.fromElements(
                new Tuple2<Long, String>(1L, "john"),
                new Tuple2<Long, String>(2L, "mick"),
                new Tuple2<Long, String>(3L, "bai"));
        // <id, age>
        DataSource<Tuple2<Long, Integer>> ages = env.fromElements(
                new Tuple2<Long, Integer>(1L, 20),
                new Tuple2<Long, Integer>(2L, 20),
                new Tuple2<Long, Integer>(4L, 40));

        // with
        names.join(ages).where(0).equalTo(0)
                .with(new JoinFunction<Tuple2<Long, String>, Tuple2<Long, Integer>,
                        Tuple3<Long, String, Integer>>() {
                    @Override
                    public Tuple3<Long, String, Integer> join(Tuple2<Long, String> t1,
                            Tuple2<Long, Integer> t2) throws Exception {
                        return new Tuple3<>(t1.f0, t1.f1, t2.f1);
                    }
                }).print();

        // map
        names.join(ages).where(0).equalTo(0)
                .map(new MapFunction<Tuple2<Tuple2<Long, String>, Tuple2<Long, Integer>>,
                        Tuple3<Long, String, Integer>>() {
                    @Override
                    public Tuple3<Long, String, Integer> map(Tuple2<Tuple2<Long, String>,
                            Tuple2<Long, Integer>> t) throws Exception {
                        return new Tuple3<>(t.f0.f0, t.f0.f1, t.f1.f1);
                    }
                }).print();
    }

    @Test
    public void testLeftOuterJoin() throws Exception {
        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();

        // <id, name>
        DataSource<Tuple2<Long, String>> names = env.fromElements(
                new Tuple2<Long, String>(1L, "john"),
                new Tuple2<Long, String>(2L, "mick"),
                new Tuple2<Long, String>(3L, "bai"));
        // <id, age>
        DataSource<Tuple2<Long, Integer>> ages = env.fromElements(
                new Tuple2<Long, Integer>(1L, 20),
                new Tuple2<Long, Integer>(2L, 40),
                new Tuple2<Long, Integer>(4L, 40));

        // leftOuterJoin
        names.leftOuterJoin(ages).where(0).equalTo(0)
                .with(new JoinFunction<Tuple2<Long, String>, Tuple2<Long, Integer>,
                        Tuple3<Long, String, Integer>>() {
                    @Override
                    public Tuple3<Long, String, Integer> join(Tuple2<Long, String> t1,
                            Tuple2<Long, Integer> t2) throws Exception {
                        if (t2 == null) {
                            return new Tuple3<>(t1.f0, t1.f1, -1);
                        }
                        return new Tuple3<>(t1.f0, t1.f1, t2.f1);
                    }
                }).print();

        // leftOuterJoin
        names.rightOuterJoin(ages).where(0).equalTo(0)
                .with(new JoinFunction<Tuple2<Long, String>, Tuple2<Long, Integer>,
                        Tuple3<Long, String, Integer>>() {
                    @Override
                    public Tuple3<Long, String, Integer> join(Tuple2<Long, String> t1,
                            Tuple2<Long, Integer> t2) throws Exception {
                        if (t1 == null) {
                            return new Tuple3<>(t2.f0, "", t2.f1);
                        }
                        return new Tuple3<>(t1.f0, t1.f1, t2.f1);
                    }
                }).print();

        // fullOuterJoin
        names.fullOuterJoin(ages).where(0).equalTo(0)
                .with(new JoinFunction<Tuple2<Long, String>, Tuple2<Long, Integer>,
                        Tuple3<Long, String, Integer>>() {
                    @Override
                    public Tuple3<Long, String, Integer> join(Tuple2<Long, String> t1,
                            Tuple2<Long, Integer> t2) throws Exception {
                        if (t1 == null) {
                            return new Tuple3<>(t2.f0, "", t2.f1);
                        } else if (t2 == null) {
                            return new Tuple3<>(t1.f0, t1.f1, -1);
                        }
                        return new Tuple3<>(t1.f0, t1.f1, t2.f1);
                    }
                }).print();
    }

    // 笛卡儿积
    @Test
    public void testCross() throws Exception {
        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();

        DataSource<String> ds1 = env.fromElements("john", "mick");
        DataSource<Integer> ds2 = env.fromElements(20, 40);

        CrossOperator.DefaultCross<String, Integer> ds3 = ds1.cross(ds2);
        ds3.print();
    }

    @Test
    public void testUnion() throws Exception {
        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();

        DataSource<String> ds1 = env.fromElements("john", "mick");
        DataSource<String> ds2 = env.fromElements("binarylei");

        UnionOperator<String> ds3 = ds1.union(ds2);
        ds3.print();
    }

    @Test
    public void testUnion2() throws Exception {
        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();

        DataSource<Tuple2<Long, String>> ds1 = env.fromElements(
                new Tuple2<Long, String>(1L, "john"),
                new Tuple2<Long, String>(1L, "mick"),
                new Tuple2<Long, String>(1L, "binarylei"),
                new Tuple2<Long, String>(2L, "a"),
                new Tuple2<Long, String>(4L, "b"),
                new Tuple2<Long, String>(3L, "c"),
                new Tuple2<Long, String>(0L, "b")
        );

        // 取前几个元素
        ds1.first(2).print();

        // 每组取2个
        System.out.println("==========");
        ds1.groupBy(0).first(2).print();

        // 组内排序
        System.out.println("==========");
        ds1.groupBy(0).sortGroup(1, Order.DESCENDING).first(2).print();

        // 全局排序，所有分区，数据量大时可能会出现问题
        System.out.println("==========");
        ds1.sortPartition(0, Order.ASCENDING)
                .sortPartition(1, Order.DESCENDING)
                .first(2).print();
    }

}
